#include "DeferredProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"

void DeferredProcess::OnInit()
{
	mUseShadows = 0;

	mCheckShadows = 0;

	mLightEntitys = mGameScene->GetEntitysByType<LocalEntity, LightMaterial>();

	float fov = mGameScene->GetCamera().GetFovY();
	float aspect = mGameScene->GetCamera().GetAspect();
	float nearZ = mGameScene->GetCamera().GetNearZ();
	float farZ = mGameScene->GetCamera().GetFarZ();

	Ubo ubo;

	ubo.viewPos = glm::vec4(mGameScene->GetCamera().GetPosition(), 0.0f) * glm::vec4(-1.0f, 1.0f, -1.0f, 1.0f);
	ubo.useShadows = mUseShadows;
	ubo.checkShadows = mCheckShadows;

	glm::mat4 shadowProj = glm::perspective(fov, aspect, nearZ, farZ);

	for (uint32 i = 0; i < mLightEntitys.size(); ++i)
	{
		LightMaterial& lightMaterial = mLightEntitys[i]->GetComponent<LightMaterial>();

		ubo.lights[i].viewMatrix = shadowProj * glm::lookAt(glm::vec3(mLightEntitys[i]->GetTransform().GetLocalPosition()), glm::vec3(lightMaterial.target), glm::vec3(0.0f, -1.0f, 0.0f));

		ubo.lights[i].position = glm::vec4(mLightEntitys[i]->GetTransform().GetLocalPosition(), 1.f);
		ubo.lights[i].target = lightMaterial.target;
		ubo.lights[i].color = lightMaterial.color;
	}




	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			RefCountPtr<RHITexture> positionRenderTarget = GetTexture("MrtProcess", "PositionRenderTarget");
			RefCountPtr<RHITexture> normalRenderTarget = GetTexture("MrtProcess", "NormalRenderTarget");
			RefCountPtr<RHITexture> albedoRenderTarget = GetTexture("MrtProcess", "AlbedoRenderTarget");
			RefCountPtr<RHITexture> depthRenderTarget = GetTexture("MrtProcess", "DepthRenderTarget");
			mCommonSampler = RHISampler::Builder().Build();

			RefCountPtr<RHITexture> shadowRenderTarget = GetTexture("ShadowProcess", "ShadowRenderTarget");
			RefCountPtr<RHISampler> shadowSampler = GetSampler("ShadowProcess", "ShadowSampler");

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<DeferredVertexShader>()
				.PushShader<DeferredFragmentShader>()
				.SetDescriptorCount(0, 1, true)
				.Build();

			mDescriptor->SetValue("samplerposition", positionRenderTarget, mCommonSampler);
			mDescriptor->SetValue("samplerNormal", normalRenderTarget, mCommonSampler);
			mDescriptor->SetValue("samplerAlbedo", albedoRenderTarget, mCommonSampler);
			mDescriptor->SetValue("samplerDepth", depthRenderTarget, mCommonSampler);
			mDescriptor->SetValue("samplerShadowMap", shadowRenderTarget, shadowSampler);
			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build(this, "ColorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetMultiple(true)
				.Build();

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void DeferredProcess::OnTick(float deltaTime)
{
	float fov = mGameScene->GetCamera().GetFovY();
	float aspect = mGameScene->GetCamera().GetAspect();
	float nearZ = mGameScene->GetCamera().GetNearZ();
	float farZ = mGameScene->GetCamera().GetFarZ();

	Ubo ubo;

	ubo.viewPos = glm::vec4(mGameScene->GetCamera().GetPosition(), 0.0f) * glm::vec4(-1.0f, 1.0f, -1.0f, 1.0f);
	ubo.useShadows = mUseShadows;
	ubo.checkShadows = mCheckShadows;

	glm::mat4 shadowProj = glm::perspective(fov, aspect, nearZ, farZ);

	for (uint32 i = 0; i < mLightEntitys.size(); ++i)
	{
		LightMaterial& lightMaterial = mLightEntitys[i]->GetComponent<LightMaterial>();

		ubo.lights[i].viewMatrix = shadowProj * glm::lookAt(glm::vec3(mLightEntitys[i]->GetTransform().GetLocalPosition()), glm::vec3(lightMaterial.target), glm::vec3(0.0f, -1.0f, 0.0f));

		ubo.lights[i].position = glm::vec4(mLightEntitys[i]->GetTransform().GetLocalPosition(), 1.f);
		ubo.lights[i].target = lightMaterial.target;
		ubo.lights[i].color = lightMaterial.color;
	}

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer->WriteMemory(&uboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void DeferredProcess::OnResize()
{
	RefCountPtr<RHITexture> positionRenderTarget = GetTexture("MrtProcess", "PositionRenderTarget");
	RefCountPtr<RHITexture> normalRenderTarget = GetTexture("MrtProcess", "NormalRenderTarget");
	RefCountPtr<RHITexture> albedoRenderTarget = GetTexture("MrtProcess", "AlbedoRenderTarget");
	RefCountPtr<RHITexture> depthRenderTarget = GetTexture("MrtProcess", "DepthRenderTarget");


	mDescriptor->SetValue("samplerposition", positionRenderTarget, mCommonSampler);
	mDescriptor->SetValue("samplerNormal", normalRenderTarget, mCommonSampler);
	mDescriptor->SetValue("samplerAlbedo", albedoRenderTarget, mCommonSampler);
	mDescriptor->SetValue("samplerDepth", depthRenderTarget, mCommonSampler);

	mDescriptor->UpdateDescriptor(mDescriptorMask);


	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build(this, "ColorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetMultiple(true)
		.Build();

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
